# 失真生成器
import os
import numpy as np
import cv2
from PIL import Image
from distortion_utils import (
    add_gaussian_noise, add_gaussian_noise_color, add_high_frequency_noise,
    add_impulse_noise, add_quantization_noise, add_gaussian_blur,
    apply_image_denoising, apply_jpeg_compression, apply_jp2k_compression,
    add_non_eccentricity_pattern_noise, apply_block_wise_distortion,
    apply_mean_shift, apply_contrast_change, apply_color_saturation,
    add_multiplicative_gaussian_noise, apply_color_quantization_dither,
    apply_color_aberrations
)

def tid2013_generator(img, dist_type, level, filename):
    """
    生成各种类型的失真图像
    
    参数:
        img: 原始图像
        dist_type: 失真类型
        level: 失真级别
        filename: 文件名
    """
    # 设置失真参数
    wn_level = [0.001, 0.005, 0.01, 0.05]  # #1 高斯噪声
    gnc_level = [0.0140, 0.0198, 0.0343, 0.0524]  # #2 彩色分量中的高斯噪声
    hfn_level = [0.001, 0.005, 0.01, 0.05]  # #5 高频噪声
    in_level = [0.005, 0.01, 0.05, 0.1]  # #6 脉冲噪声
    qn_level = [255//27, 255//39, 255//55, 255//76]  # #7 量化噪声
    gblur_level = [7, 15, 39, 91]  # #8 高斯模糊
    id_level = [0.001, 0.005, 0.01, 0.05]  # #9 图像去噪
    jpeg_level = [43, 12, 7, 4]  # #10 JPEG压缩
    jp2k_level = [0.46, 0.16, 0.07, 0.04]  # #11 JP2K压缩
    nepn_level = [30, 70, 150, 300]  # #14 非偏心模式噪声
    bw_level = [2, 4, 8, 16, 32]  # #15 不同强度的局部块状失真
    ms_level = [15, 30, 45, 60]  # #16 均值偏移
    cc_level = [0.85, 0.7, 0.55, 0.4]  # #17 对比度变化
    cs_level = [0.4, 0, -0.4, -0.8]  # #18 色彩饱和度
    mgn_level = [0.05, 0.09, 0.13, 0.2]  # #19 乘性高斯噪声
    cqd_level = [64, 32, 16, 8, 4]  # #22 颜色量化抖动
    ca_level = [2, 6, 10, 14]  # #23 色差

    # 获取不带扩展名的文件名
    base_filename = os.path.splitext(filename)[0]
    
    # 根据失真类型生成失真图像
    if dist_type == 1:  # 高斯噪声
        strname = './GN/GN'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_gaussian_noise(img, wn_level[level-1])
        
    elif dist_type == 2:  # 彩色分量中的高斯噪声
        strname = './GNC/GNC'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_gaussian_noise_color(img, gnc_level[level-1])
        
    elif dist_type == 5:  # 高频噪声
        strname = './HFN/HFN'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_high_frequency_noise(img, hfn_level[level-1])
        
    elif dist_type == 6:  # 脉冲噪声
        strname = './IN/IN'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_impulse_noise(img, in_level[level-1])
        
    elif dist_type == 7:  # 量化噪声
        strname = './QN/QN'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_quantization_noise(img, qn_level[level-1])
        
    elif dist_type == 8:  # 高斯模糊
        strname = './GB/GB'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_gaussian_blur(img, gblur_level[level-1])
        
    elif dist_type == 9:  # 图像去噪
        strname = './ID/ID'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = apply_image_denoising(img, id_level[level-1])
        
    elif dist_type == 10:  # JPEG压缩
        strname = './JPEG/JPEG'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.jpg")
        # 创建目录
        os.makedirs(os.path.dirname(test_name), exist_ok=True)
        # 直接保存为JPEG格式
        apply_jpeg_compression(img, test_name, jpeg_level[level-1])
        return  # 直接返回，因为已经保存了图像
        
    elif dist_type == 11:  # JP2K压缩
        strname = './JP2K/JP2K'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.jp2")
        # 创建目录
        os.makedirs(os.path.dirname(test_name), exist_ok=True)
        # 直接保存为JP2K格式
        apply_jp2k_compression(img, test_name, 24 / jp2k_level[level-1])
        return  # 直接返回，因为已经保存了图像
        
    elif dist_type == 14:  # 非偏心模式噪声
        strname = './NEPN/NEPN'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_non_eccentricity_pattern_noise(img, nepn_level[level-1])
        
    elif dist_type == 15:  # 不同强度的局部块状失真
        strname = './BW/BW'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = apply_block_wise_distortion(img, bw_level[level-1], level)
        
    elif dist_type == 16:  # 均值偏移
        strname = './MSH/MSH'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = apply_mean_shift(img, ms_level[level-1])
        
    elif dist_type == 17:  # 对比度变化
        strname = './CCL/CCL'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = apply_contrast_change(img, cc_level[level-1])
        
    elif dist_type == 18:  # 色彩饱和度
        strname = './CS/CS'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = apply_color_saturation(img, cs_level[level-1])
        
    elif dist_type == 19:  # 乘性高斯噪声
        strname = './MGN/MGN'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = add_multiplicative_gaussian_noise(img, mgn_level[level-1])
        
    elif dist_type == 22:  # 颜色量化抖动
        strname = './CQD/CQD'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = apply_color_quantization_dither(img, cqd_level[level-1])
        
    elif dist_type == 23:  # 色差
        strname = './CA/CA'
        test_name = os.path.join(f"{strname}{level}", f"{base_filename}.bmp")
        distorted_img = apply_color_aberrations(img, ca_level[level-1])
    
    else:
        print(f"未知的失真类型: {dist_type}")
        return
    
    # 创建目录并保存图像
    os.makedirs(os.path.dirname(test_name), exist_ok=True)
    cv2.imwrite(test_name, distorted_img)